import json
import numpy as np
import matplotlib.pyplot as plt

# Load JSON data
with open("generative_pm/results_Llama3_length1shot.json", 'r') as file:
    data = json.load(file)

chosen_score_list = data["chosen_score_list"]

# Convert list to numpy array
scores = np.array(chosen_score_list)

# Calculate absolute distances from 0.5
distances = np.abs(scores - 0.5)

# Define bins with a width of 0.1 between 0 and 0.5
bins = np.arange(0, 0.6, 0.1)

# Bin each score into the corresponding range
bin_indices = np.digitize(distances, bins, right=True) - 1

# Initialize dictionary to store bin results
bin_results = {}

# Iterate through bins
for i in range(len(bins) - 1):
    bin_scores = scores[bin_indices == i]
    greater_than_half = np.sum(bin_scores > 0.5)
    proportion_greater = greater_than_half / len(bin_scores) if bin_scores.size > 0 else 0
    bin_results[(bins[i], bins[i + 1])] = {
        'greater_than_half': proportion_greater,
        'sample_size': len(bin_scores)
    }

# Prepare data for plotting
bin_ranges = [(f"{k[0]+0.5:.1f}-{k[1]+0.5:.1f}", v['greater_than_half'], v['sample_size']) for k, v in bin_results.items() if k[0] <= 0.5]

# Create figure and axis
fig, ax1 = plt.subplots(figsize=(8, 6))
bar_width = 0.7
index = np.arange(len(bin_ranges))

# Plot bar chart for 'greater than 0.5'
bars1 = ax1.bar(index, [x[1] for x in bin_ranges], bar_width, color='royalblue', alpha=0.7, label='Agreement')

# 在柱状图上标注数值
for bar in bars1:
    yval = bar.get_height()
    ax1.text(bar.get_x() + bar.get_width() / 2, yval, round(yval, 2), ha='center', va='bottom', fontsize=26)



# Set labels and title
ax1.set_xlabel('Preference Strength', fontsize=26)
ax1.set_ylabel('Agreement', fontsize=26)
# ax1.set_title('Ultra Feedback', fontsize=26)

# Customize x-axis
ax1.set_xticks(index)
ax1.set_xticklabels([f"{r[0]}" for r in bin_ranges], fontsize=26)

# Remove top and right spines
ax1.spines['top'].set_visible(False)
ax1.spines['right'].set_visible(False)

# Set tick parameters
ax1.tick_params(axis='both', labelsize=18)


# Add legend
ax1.legend(loc='upper right', fontsize=22)

proportion_list = [230/(230+142+105+140+374),142/(230+142+105+140+374),105/(230+142+105+140+374),
                   140/(230+142+105+140+374),374/(230+142+105+140+374)]
# 创建第二个Y轴
ax2 = ax1.twinx()
# 绘制proportion_list的折线图
line2, = ax2.plot(range(len(proportion_list)), proportion_list, color='orange', marker='o', markersize=10, linewidth=4, label='Proportion')
# 设置第二个Y轴的标签
ax2.set_ylabel('Proportion', fontsize=26)
# 美化折线图标记
ax2.tick_params(axis='y', labelsize=22)

# 添加图例
ax1.legend(loc='upper left', fontsize=22)
ax2.legend(loc='upper left', fontsize=22, bbox_to_anchor=(0, 0.88))

ax1.set_ylim(0.46, 0.83)
ax2.set_ylim(0, 0.6)

# Save and show the plot
plt.tight_layout()
plt.savefig('generative_pm/proportion_bins-Llama3_length1shot.png')
plt.show()
